from openvino.tools import mo
from openvino.runtime import serialize

class OpenVINOExporter:
    def __init__(self, onnx_path):
        self.onnx_path = onnx_path
        
    def convert_to_ir(self, output_dir):
        # 模型优化配置
        config = {
            'input_shape': [1,3,640,640],
            'mean_values': '[0,0,0]',
            'scale_values': '[255,255,255]',
            'compress_to_fp16': True
        }
        
        # 执行模型优化
        ov_model = mo.convert_model(
            self.onnx_path,
            **config
        )
        
        # 保存IR文件
        serialize(ov_model, f'{output_dir}/yolov8.xml')
        print(f'Successfully exported OpenVINO IR to {output_dir}')

if __name__ == '__main__':
    exporter = OpenVINOExporter('yolov8s.onnx')
    exporter.convert_to_ir('../runtime/openvino_engine')